/**
 * @file encode_dlm_cls_v2_o.cpp
 * @brief ENNX分类模型文件编码和转换功能的实现文件
 * @details 该文件实现了用于处理ENNX分类模型文件的编码和转换相关功能:
 *          - 支持模型文件的加密和解密
 *          - 支持模型格式转换(ENNX到TensorRT)
 *          - 支持动态/静态batch处理
 *          - 支持不同精度模式(FP32/FP16/INT8)
 * @author 
 * @date 2024-01
 */
#include "encrypt_rec1_o.h"  // Include YOLO v6 class definition
#include "encode_dlm_rec_v1_o.h"

#include <thread>  // Thread library
#include <chrono>  // Time library
#define WIN32_LEAN_AND_MEAN // Avoid including WinSock.h
#include <WinSock2.h>
#include <Windows.h>

#include <iostream>  // Standard I/O streams
#include <cstdio>    // Standard I/O
#include <direct.h>   // Directory operations
#include <vector>     // Dynamic array (vector)
#include "boost/filesystem/operations.hpp"  // Boost filesystem for file and directory operations
#include "boost/filesystem/path.hpp"        // Boost filesystem for path operations
#include "boost/filesystem/operations.hpp"  // Duplicate include, may be for emphasis
#include "boost/filesystem/fstream.hpp"      // Boost filesystem for file streams

#include "common/device_utility.h"

/**
 * @brief 枚举指定路径下具有特定扩展名的所有文件
 * @param path 要枚举的目录路径
 * @param extension 要匹配的文件扩展名
 * @return 返回包含所有匹配文件路径的向量
 */
std::vector<std::string> EnumFiles(
     const std::string& path,
     const std::string &extension)
{
    namespace fs = boost::filesystem;
    fs::path _path(path);
    fs::path fullpath = fs::system_complete(_path);

    std::vector<std::string> ret;

    if (!fs::exists(fullpath) && !fs::is_directory(fullpath)) {
        return ret;
    }
    fs::directory_iterator end_iter;
    for (fs::directory_iterator iter(fullpath); iter != end_iter; iter++) {
        try {
            if (fs::is_directory(*iter)) {
            }
            else if (iter->path().extension() == extension) {
                ret.push_back(iter->path().string());
            }
        } catch (const std::exception&) {
            continue;
        }
    }
    return ret;
}

/**
 * 枚举模型配置文件并执行模型转换
 * 
 * @param mon_folder 模型文件夹路径
 * @param dest_folder 目标文件夹路径
 * @param mark_info 用于生成目标文件名的标记信息
 */
void enum_models(std::string& mon_folder, std::string& dest_folder, std::string & mark_info) {
    namespace fs = boost::filesystem;
    std::string gpu_model = GetGPUModel();
    std::cout <<"gpu_model : " << gpu_model << std::endl;
    if (gpu_model.empty()) {
        std::cout << "Unknown Gpu Model! \n";
        return;
    }

    fs::path _dest_folder(dest_folder);
    if (!fs::exists(_dest_folder)) {
        return;
    }
    std::cout << _dest_folder.string() <<" exist"<< std::endl;

    std::vector<std::string> cfg_files;
    cfg_files = EnumFiles(
        mon_folder,
        ".yaml");
    std::cout << "cfg_files : " << cfg_files.size() << std::endl;

    for (const auto & cfg_file : cfg_files) {
        std::cout << "This cfg : " << cfg_file.c_str() << std::endl;
        st_trans_model_input input;
        fs::path _cfg_path(cfg_file);
        input.cfg_path = _cfg_path.string();
        std::cout << "input.cfg_path : " << input.cfg_path << std::endl;

        fs::path _model_path;
        {
            fs::path _path = _cfg_path.replace_extension(".ennx");
            if (fs::exists(_path)) {
                _model_path = _path;
            }
        }
        {
            fs::path _path = _cfg_path.replace_extension(".onnx");
            if (fs::exists(_path)) {
                _model_path = _path;
                std::cout << "model onnx: " << _model_path.string() << std::endl;
            }
        }
        if (_model_path.empty()) {
            continue;
        }

        fs::path _gdats_path(dest_folder
            + "\\" + _cfg_path.stem().string() +
            mark_info +
			// 添加文件扩展名
            FILE_EXTENSION );
        std::cout << "gdats_path : " << _gdats_path.string() << std::endl;
        if (fs::exists(_gdats_path)) {
            continue;
        }

        std::cout << "START===============================================\n";
        input.model_path = _model_path.string();
        input.gpu_model = gpu_model;
        input.trans_path = _gdats_path.string() + "-tmp";

		// 创建模型转换器实例
        modeltrans_segv2 modelConverter;
        if (modelConverter.TransModel(input)) {
            std::cout << "DONE================================================\n";
            int renameResult = std::rename(
                input.trans_path.c_str(),
                _gdats_path.string().c_str());
            if (renameResult == 0) {
                std::cout << "Conversion successful and file has been renamed.\n";
            } else {
                std::cout << "Conversion successful, but file renaming failed.\n";
            }
        } else {
            std::cout << "FAIL!================================================\n";
        }
    }
}

int main(int argc, char* argv[]) {
    
    try {
        // 获取当前工作目录
        char cur_folder[260];
        _getcwd(cur_folder, sizeof(cur_folder));
        printf("Current working directory: \n %s\n", cur_folder);

        std::string cwd = cur_folder;

        // 初始化路径
        std::string mon_folder = cwd + "/model";
        std::string dest_folder = cwd;
        std::string mark_info = "v1";

        // 解析命令行参数
        if (argc >= 2) {
            mon_folder = argv[1];    // 第一个参数: 模型文件所在的源文件夹路径
        }
        if (argc >= 3) {
            dest_folder = argv[2];   // 第二个参数: 转换后模型文件的目标存储路径
        }
        if (argc >= 4) {
            mark_info = argv[3];     // 第三个参数: 模型版本标记,用于区分不同版本
        }

        // 执行模型转换
        enum_models(mon_folder, dest_folder, mark_info);

        printf("Done!\n");
    }
    catch (...) {
        printf("Exception!\n");
    }

    Sleep(3000);
    return 0;
}
